package com.zrkizzy.common.limit.strategy;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import com.zrkizzy.common.core.utils.StringUtil;
import com.zrkizzy.common.limit.LimitManager;
import com.zrkizzy.common.limit.annotation.AccessLimit;
import com.zrkizzy.common.limit.exception.LimitErrorCode;
import lombok.extern.slf4j.Slf4j;

import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * Guava限流策略
 *
 * @author zhangrongkang
 * @since 2024/2/14
 */
@Slf4j
public class GuavaLimit implements LimitManager {

    /**
     * 不同的接口进行不同的流量控制
     * <p>
     *     key为 AccessLimit.key()
     * </p>
     *
     */
    private final Map<String, RateLimiter> accessLimitMap = Maps.newConcurrentMap();

    /**
     * 请求限流
     *
     * @param accessLimit 自定义限流注解
     * @return 是否请求成功
     */
    @Override
    public boolean tryAccess(AccessLimit accessLimit) {
        // 获取限流参数，进入当前方法时说明AccessLimit一定不为空
        String key = accessLimit.key();
        // 校验Key是否合法
        if (StringUtil.isBlank(key)) {
            throw LimitErrorCode.INTERFACE_KEY_NOT_NULL.exception();
        }
        int max = accessLimit.max();
        long timeout = accessLimit.timeout();
        TimeUnit timeunit = accessLimit.timeunit();
        // Guava令牌桶算法限流类
        RateLimiter rateLimiter = null;
        // 验证缓存中是否有命中的Key
        if (!accessLimitMap.containsKey(key)) {
            // 创建令牌桶
            rateLimiter = RateLimiter.create(max);
            accessLimitMap.put(key, rateLimiter);
            log.info("新建令牌桶, Key: {}, 容量: {}", key, max);
        }
        // 获取令牌
        rateLimiter = accessLimitMap.get(key);
        // 执行限流方法
        return rateLimiter.tryAcquire(timeout, timeunit);
    }

}
